import h5py
import os
import numpy as np
import SimpleITK as sitk
from tqdm import tqdm
# 四种模态的mri图像
modalities = ('flair', 't1ce', 't1', 't2')

# train
train_set = {
        'input_data_path': '/root/autodl-tmp/BrainTumorSegmentation1/data_set/BraTS2021/data',  # 四个模态数据所在地址
        'output_data_path': '/root/autodl-tmp/BrainTumorSegmentation1/data_set/BraTS2021/dataset',  # 预处理数据的输出地址
        'file_list_txt': 'train.txt',  # 训练集名单（有标签）
        }


def process_h5(path, out_path):
    """
        将数据保存为dtype=float32。
        使用z-score进行标准化，但保持背景值为零
    """
    # SimpleITK读取图像默认是是 DxHxW，这里转为 HxWxD
    label = sitk.GetArrayFromImage(sitk.ReadImage(path + 'seg.nii.gz')).transpose(1,2,0)  # [240, 240, 155]
    # print(label.shape)
    # 堆叠四种模态的图像，4 x (H,W,D) -> (4,H,W,D) 
    images = np.stack([sitk.GetArrayFromImage(sitk.ReadImage(path + modal + '.nii.gz')).transpose(1,2,0) for modal in modalities], 0)  # [240,240,155]

    label = label.astype(np.uint8)
    images = images.astype(np.float32)
    # print(images.shape)

    case_name = os.path.basename(path)  # 从给定路径中获取基本文件名（即不包括目录路径的文件名）

    # path = os.path.join(out_path,case_name) # 将输出路径 out_path 和文件名 case_name 组合成完整的输出文件路径。
    # print('path=',path)
    # output = path + 'mri_norm2.h5'
    #output_dir = os.path.join(out_path, case_name)  # Output directory path
    #os.makedirs(output_dir, exist_ok=True)  # Create the directory if it doesn't exist
    # output = os.path.join(path, 'mri_norm2.h5') # 保存完整的输出文件路径，用于保存处理后的数据
    output = os.path.join(out_path, case_name + 'mri_norm2.h5') # 将输出路径 out_path 和文件名 case_name + 'mri_norm2.h5' 组合成一个完整的输出文件路径
    # print('output=',output)
    # 对第一个通道求和，如果四个模态都为0，则标记为背景(False)
    mask = images.sum(0) > 0
    for k in range(4):

        x = images[k,...]  # 提取了images数组中第k个模态的数据并将其存储在变量x中
        y = x[mask] # 筛选出x中的非背景区域

        # 对背景外的区域进行归一化
        x[mask] = x[mask] - y.mean()
        x[mask] = x[mask] / y.std()

        images[k,...] = x # 将经过归一化处理后的图像数据 x 赋值回原始的 images 数组中的第 k 个模态
    # print('case_name=,images.shape=,label.shape=',case_name,images.shape,label.shape)


    f = h5py.File(output, 'w') # 创建了一个新的 HDF5 文件对象 f，以写入（'w'）模式打开了一个文件，文件名由变量 output 指定
    # 在f中创建两个数据集，image和label，并用gzip压缩
    f.create_dataset('image', data=images, compression="gzip")
    f.create_dataset('label', data=label, compression="gzip")
    f.close()


def process_dataset(dset):
    root, out_path = dset['input_data_path'], dset['output_data_path']
    file_list = os.path.join(root, dset['file_list_txt'])
    # print(file_list)
    subjects = open(file_list).read().splitlines()
    names = ['BraTS2021_' + sub for sub in subjects]
    paths = [os.path.join(root, name, name + '_') for name in names]
    # print(paths)

    for path in tqdm(paths):
        process_h5(path, out_path)
        # break
    print('Done')


if __name__ == '__main__':
    process_dataset(train_set)
